# src/exphd.py

import os
import re
import logging
from collections import defaultdict

import numpy as np
import pandas as pd
import matplotlib

matplotlib.use("Agg")  # Set matplotlib to use non-interactive backend
import matplotlib.pyplot as plt

logger = logging.getLogger(__name__)


class ExpHandler:
    def __init__(self):
        self.groups = {
            "Gibbs energy (G)": ["G-T.EXP", "GI-T.EXP"],
            "Enthalpy (H)": ["H-T.EXP", "HI-T.EXP"],
            "Entropy (S)": ["S-T.EXP", "SI-T.EXP"],
            "Amount of phase (NP)": ["NP-T.EXP", "XBL-T.EXP"],
            "Composistion of phase (X)": ["X-T.EXP", "XCF-T.EXP"],
            "Site fraction (Y)": ["Y-T.EXP", "YSITE-T.EXP"],
        }

    def split_exp(self, exp_file: str):
        """Split EXP file into blocks starting with '$ PLOTTED COLUMNS ARE'.

        Args:
            exp_file: EXP file

        Yields:
            list[str]: block
        """
        with open(exp_file, "r", encoding="utf-8") as f:
            lines = f.readlines()

        block = []
        for line in lines:
            if "$ PLOTTED COLUMNS ARE" in line and len(block) > 0:
                yield block
                block = []
            block.append(line)

        # Return the last block
        if len(block) > 0:
            yield block

    def block2data(self, block: list[str]) -> pd.DataFrame:
        """Convert block to DataFrame

        Args:
            block: block

        Returns:
            pd.DataFrame: Converted DataFrame
        """
        # Match independent variable T and dependent variables
        match = re.search(r": (\S+) and (\S+)", block[0])
        if not match:
            return pd.DataFrame()
        column1, column2 = match.group(1), match.group(2)

        column1_data = []
        column2_data = []
        flag = False
        for line in block[1:]:
            parts = line.split()
            if len(parts) < 2:
                continue
            # Find the first line of data
            if "M" == parts[-1]:
                flag = True
            if flag:
                try:
                    column1_data.append(float(parts[0]))
                    column2_data.append(float(parts[1]))
                except ValueError:
                    continue

        if len(column1_data) > 0 and len(column2_data) > 0:
            return pd.DataFrame({column1: column1_data, column2: column2_data})

    def handle_exps(self, exp_files: list[str]) -> pd.DataFrame:
        """Process EXP files of the same type

        Args:
            exp_files: Multiple EXP files

        Returns:
            pd.DataFrame: Processed DataFrame
        """

        # Group block df by column names
        exp_groups = defaultdict(list)
        for exp_file in exp_files:
            for block in self.split_exp(exp_file):
                df = self.block2data(block)
                if df.empty or len(df) < 2:
                    continue
                col1, col2 = df.columns.str.strip()
                df.dropna(subset=[col1, col2], inplace=True)
                exp_groups[col2].append(df[[col1, col2]])

        if not exp_groups:
            return pd.DataFrame()

        # Concatenate data of the same type vertically
        series_dict = defaultdict(list)
        for col2, dfs in exp_groups.items():
            combined = pd.concat(dfs, ignore_index=True)
            col1 = combined.columns[0]
            combined = combined.drop_duplicates(subset=[col1], keep="first")
            series_dict[col2] = combined.set_index(col1)[col2]

        # Merge all columns horizontally
        exp_df = pd.concat(series_dict, axis=1)
        exp_df.sort_index(inplace=True)
        exp_df = exp_df[~exp_df.index.duplicated(keep="first")]

        return exp_df

    def plot_df(self, df: pd.DataFrame, name: str, out_dir: str):
        """Plot DataFrame

        Args:
            df: DataFrame
            name: Chart name to plot
            out_dir: Output dir
        """
        # Find appropriate submarker
        all_markers = ["o", "s", "^", "D", "*", "P", "H", "X"]
        num_col = len(df.columns)
        if num_col > len(all_markers):
            markers = all_markers[:num_col]
        else:
            for i in range(len(all_markers), 0, -1):
                if num_col % i == 0:
                    markers = all_markers * (num_col // i)
                    break
            else:
                markers = all_markers[:num_col]
        num_mk = len(markers)
        mark_every = 100  # Plot a point every 100 data points
        plt.figure(figsize=(12, 8))
        for i, column in enumerate(df.columns):
            plt.plot(
                df.index,
                df[column],
                label=column,
                marker=markers[i % num_mk],
                markevery=mark_every,
            )
        plt.title(f"T-dependency of {name}", fontsize=20)
        plt.xlabel("Temperature, K", fontsize=20)
        plt.ylabel(name, fontsize=20)
        plt.grid()
        plt.legend(loc="best", bbox_to_anchor=(1, 1), fontsize=10)
        plt.tight_layout(rect=(0, 0, 1, 0.95))
        plt.savefig(os.path.join(out_dir, f"T-dep-{name}.png"), bbox_inches="tight")
        plt.close()

    def exps2xlsx(self, exp_dir: str, out_file: str = "tc_exps.xlsx"):
        """Convert multiple EXP files in dir to XLSX file

        Args: str: EXP file dir
              str: Output file dir

        Returns: str: Output filename
        """
        dfs: dict[str, pd.DataFrame] = {}
        for name, group in self.groups.items():
            exp_files = [os.path.join(exp_dir, file) for file in group]
            exp_files = [f for f in exp_files if os.path.isfile(f)]
            df = self.handle_exps(exp_files)
            if df.empty:
                logger.warning(f"No valid EXP data for {name} in {exp_dir}")
                continue

            dfs[name] = df

        if len(dfs) <= 0:
            logger.warning(f"No valid EXP data in {exp_dir}")
            return ""

        output = os.path.join(exp_dir, out_file)
        logger.info(f"Found {len(dfs)} EXP data in {exp_dir}, save to {output}")
        with pd.ExcelWriter(output) as writer:
            for name, df in dfs.items():
                df.to_excel(writer, sheet_name=name)
                self.plot_df(df, name, exp_dir)
